// 从迷宫(PNG图片，白色背景，黑色线条)中寻找正确路径，
// 入口用红色(0xFF0000)的点标记，
// 出口用蓝色(0x0000FF)的点标记
package main

import (
	"flag"
	"image"
	"image/color"
	"image/png"
	"os"
)

func getPngImage(fn string) *image.NRGBA {
	fp, err := os.Open(fn)
	if err != nil {
		panic(err)
	}
	defer fp.Close()
	img, err := png.Decode(fp)
	if err != nil {
		panic(err)
	}
	rect := img.Bounds()
	w := rect.Max.X - rect.Min.X
	h := rect.Max.Y - rect.Min.Y
	img1 := image.NewNRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			img1.Set(x, y, img.At(rect.Min.X+x, rect.Min.Y+y))
		}
	}
	return img1
}

type Point struct {
	X, Y int
	Last int
}

var points []Point

func getNextPoints(n int) []Point {
	res := []Point{}
	for i := -1; i <= 1; i++ {
		for j := -1; j <= 1; j++ {
			if i == 0 && j == 0 {
				continue
			}
			res = append(res, Point{X: points[n].X + i, Y: points[n].Y + j, Last: n})
		}
	}
	return res
}

func saveImage(img image.Image, fn string) error {
	fpout, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer fpout.Close()
	err = png.Encode(fpout, img)
	return err
}

func main() {
	var fn = flag.String("png", "", "PNG image path")
	var m = flag.Int("mode", 0, "模式：0-实线，1-图纸")
	var outFn = flag.String("o", "out.png", "output PNG file name")
	flag.Parse()
	img := getPngImage(*fn)
	switch *m {
	case 0:
		findWaySolid(img)
	case 1:
		findWayPipe(img)
	}

	saveImage(img, *outFn)
}

func findWaySolid(img *image.NRGBA) {
	x, y := findStart(img)
	points = []Point{{X: x, Y: y, Last: -1}}
	var exits = []int{}
	var n = 0
	c := color.RGBA{R: 166, G: 166, B: 166, A: 255}
	for {
		if n >= len(points) {
			break
		}
		p := points[n]
		img.Set(p.X, p.Y, c)
		for _, p1 := range getNextPoints(n) {
			if colorBlue(img, p1.X, p1.Y) {
				exits = append(exits, n)
				break
			}
			c1 := img.At(p1.X, p1.Y)
			if !equalColor(c1, c) && !colorWhite(img, p1.X, p1.Y) {
				img.Set(p1.X, p1.Y, c)
				points = append(points, p1)
			}
		}
		n++
	}

	for _, i := range exits {
		drawPath(img, i)
	}
}

func findWayPipe(img *image.NRGBA) {
	x, y := findGreen(img)
	points = []Point{{X: x, Y: y, Last: -1}}
	var exits = []int{}
	var entrance int
	var n = 0
	c := color.RGBA{R: 166, G: 166, B: 166, A: 255}
	ce := color.RGBA{R: 255, G: 0, B: 0, A: 255}
	for {
		if n >= len(points) {
			break
		}
		p := points[n]

		img.Set(p.X, p.Y, c)

		for _, p1 := range getNextPoints(n) {
			if colorBlue(img, p1.X, p1.Y) {
				exits = append(exits, n)
				break
			}

			c1 := img.At(p1.X, p1.Y)

			if equalColor(c1, ce) {
				entrance = n
			}

			if colorWhite(img, p1.X, p1.Y) {
				img.Set(p1.X, p1.Y, c)
				points = append(points, p1)
			}
		}
		n++
	}

	for _, i := range exits {
		drawPath(img, i)
	}
	drawPath(img, entrance)
}

func equalColor(c0 color.Color, c1 color.RGBA) bool {
	r, g, b, _ := c0.RGBA()
	r = r >> 8
	g = g >> 8
	b = b >> 8

	if c1.R == uint8(r) && c1.G == uint8(g) && c1.B == uint8(b) {
		return true
	}
	return false
}

func drawPath(img *image.NRGBA, n int) {
	c := color.RGBA{R: 255, G: 0, B: 0, A: 255}
	for {
		if n == -1 {
			break
		}
		img.Set(points[n].X, points[n].Y, c)
		n = points[n].Last
	}
}

func colorBlue(img image.Image, x, y int) bool {
	r, g, b, _ := img.At(x, y).RGBA()
	r = r >> 8
	g = g >> 8
	b = b >> 8
	if r < 50 && g < 50 && b > 200 {
		return true
	}
	return false
}

func findStart(img image.Image) (x, y int) {
	rect := img.Bounds()
	var xall, yall int
	var n int
	for i := rect.Min.X; i < rect.Max.X; i++ {
		for j := rect.Min.Y; j < rect.Max.Y; j++ {
			r, g, b, _ := img.At(i, j).RGBA()
			if r > 200 && g < 50 && b < 50 {
				xall += i
				yall += j
				n++
			}
		}
	}
	x = xall / n
	y = yall / n
	return
}

func findGreen(img *image.NRGBA) (x, y int) {
	rect := img.Bounds()
	var xall, yall int
	var n int
	for i := rect.Min.X; i < rect.Max.X; i++ {
		for j := rect.Min.Y; j < rect.Max.Y; j++ {
			r, g, b, _ := img.At(i, j).RGBA()
			if g > 200 && r < 50 && b < 50 {
				xall += i
				yall += j
				n++
				img.Set(i, j, color.White)
			}
		}
	}
	x = xall / n
	y = yall / n
	return
}

func colorWhite(img image.Image, x, y int) bool {
	r, g, b, _ := img.At(x, y).RGBA()
	r = r >> 8
	g = g >> 8
	b = b >> 8
	if r > 240 && g > 240 && b > 240 {
		return true
	}
	return false
}
